/*
 * Copyright (c) 2022 HiSilicon (Shanghai) Technologies CO., LIMITED.
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

/* -------------------------------------------------------------------------- */
/* --- PUBLIC DEPENDANCIES -------------------------------------------------- */
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <errno.h>
#include <fcntl.h>
#include <sys/mman.h>
#include <sys/stat.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <unistd.h>
#include <signal.h>
#include <pthread.h>
#include <sys/prctl.h>
#include <math.h>
#include <pthread.h>

#include "hi_common.h"
#include "hi_comm_sys.h"
#include "hi_comm_svp.h"
#include "sample_comm.h"

/* -------------------------------------------------------------------------- */
/* --- COMMON DEPENDANCIES -------------------------------------------------- */
#include "sample_comm_svp.h"
#include "sample_comm_nnie.h"
#include "sample_comm_ive.h"

#include "trace.h"
#include "ffmpeg_rtmp.h"

/* -------------------------------------------------------------------------- */
/* --- PRIVATE MARCOS ------------------------------------------------------- */
static AVFormatContext *output_format_context = NULL;
static AVStream *video_stream = NULL, *audio_stream = NULL;
static const char *rtmp_url = "rtmp://192.168.200.100/live"; 
static AVRational s_tTime_base = {1, 30};
static AVRational us_time_base = {1, 1000000};

static tH264SPSPPSHeader h264_sps_pps_header;

static pthread_mutex_t write_frame_mutex;
/* function : Init ffmpeg network
 * create tcp connect 
 * init frame data buf queue
 */
HI_S32 initFFmpegNetwork( HI_VOID )
{
    HI_S32 s32Ret = HI_SUCCESS;
    /* step 1: register encoder and format */
    av_register_all();
    s32Ret = avformat_network_init();
    MSG_DEBUG(DEBUG_FFMPEG, "[INFO] : [initFFmpegNetwork] register encoder and format success, ret = %d.\n", s32Ret);

    /* step 2: creat rtmp output flv context */
    s32Ret = avformat_alloc_output_context2(&output_format_context, NULL, "flv", rtmp_url);
    if (s32Ret < 0)
    {
        MSG_DEBUG(DEBUG_FFMPEG, "[ERROR] : [initFFmpegNetwork] avformat_alloc_output_context2 failed.\n");
        goto FAIL_0;
    }

    /* step 3: creat new video frame stream */
    video_stream = avformat_new_stream(output_format_context, NULL);
    if (!video_stream)
    {
        MSG_DEBUG(DEBUG_FFMPEG, "[ERROR] : [initFFmpegNetwork] avformat_new_video_stream failed.\n");
        goto FAIL_0;
    }

    // video_stream->time_base = s_tTime_base;
    video_stream->nb_frames = 0;

    video_stream->codecpar->codec_tag = 0;
    video_stream->codecpar->codec_type = AVMEDIA_TYPE_VIDEO;
    video_stream->codecpar->codec_id = AV_CODEC_ID_H264;
    video_stream->codecpar->width = 1920;
    video_stream->codecpar->height = 1080;
    video_stream->codecpar->format = AV_PIX_FMT_YUV420P;
    video_stream->codecpar->bit_rate = 2500 * 1024;

    // create audio_stream
    audio_stream = avformat_new_stream(output_format_context, NULL);
    if (!audio_stream)
    {
        MSG_DEBUG(DEBUG_FFMPEG, "[ERROR] : [initFFmpegNetwork] avformat_new_audio_stream failed.\n");
        goto FAIL_0;
    }
    
    // audio_stream->codecpar->codec_tag = 1;
    audio_stream->codecpar->codec_type = AVMEDIA_TYPE_AUDIO;
    audio_stream->codecpar->codec_id = AV_CODEC_ID_AAC;
    audio_stream->codecpar->sample_rate = 48000;
    audio_stream->codecpar->channels = 1;
    audio_stream->codecpar->channel_layout = AV_CH_LAYOUT_MONO;
    audio_stream->codecpar->format = AV_SAMPLE_FMT_S16;
    audio_stream->codecpar->bit_rate = 128000;

    pthread_mutex_init(&(write_frame_mutex), NULL);

    return s32Ret;
    
FAIL_0:
    if (output_format_context)
    {
        if (!(output_format_context->oformat->flags & AVFMT_NOFILE))
        {
            avio_closep(&output_format_context->pb);
        }
        avformat_free_context(output_format_context);
    }
    return HI_FAILURE;
}

HI_VOID cleanFFmpegSetting( HI_VOID )
{
    if (output_format_context)
    {
        /* step 1: write format context tailer */
        av_write_trailer(output_format_context);
        
        /* step 2: clean format context */
        if (!(output_format_context->oformat->flags & AVFMT_NOFILE))
        {
            avio_closep(&output_format_context->pb);
        }

        /* step 3: free format context */
        avformat_free_context(output_format_context);
        MSG_DEBUG(DEBUG_FFMPEG, "[INFO] : [cleanFFmpegSetting] clean up format context information.\n");
        
        pthread_mutex_destroy(&(write_frame_mutex));
    }
}

HI_VOID init_h264_sps_pps_header( int size )
{
    /* malloc h264_sps_pps_header */
    h264_sps_pps_header.sps_pps_buf = malloc(size);
    h264_sps_pps_header.sps_pps_size = 0;
    h264_sps_pps_header.has_sps_pps = HI_FALSE;
    h264_sps_pps_header.isStreamInit = HI_FALSE;
    h264_sps_pps_header.isKeyFrame = HI_FALSE;
    h264_sps_pps_header.isFirstIdrFrame = HI_FALSE;
    h264_sps_pps_header.start_pts = 0;
    MSG_DEBUG(DEBUG_FFMPEG, "[INFO] : [initFrameQueueMalloc] malloc memeory success.\n");
}

HI_VOID freeH264SpsPpsHeader( HI_VOID )
{
    free(h264_sps_pps_header.sps_pps_buf);
    h264_sps_pps_header.sps_pps_size = 0;
    h264_sps_pps_header.has_sps_pps = HI_FALSE;
    h264_sps_pps_header.isStreamInit = HI_FALSE;
    h264_sps_pps_header.isKeyFrame = HI_FALSE;
    h264_sps_pps_header.isFirstIdrFrame = HI_FALSE;
    h264_sps_pps_header.start_pts = 0;
}

static HI_S32 sendH264Frame(uint8_t* data, HI_U32 size, HI_BOOL isKey, uint64_t pts_)
{
    HI_S32 s32Ret = HI_SUCCESS;

    /* step 1: init h264 packet */
    AVPacket *pkt = av_packet_alloc();
    if (!pkt)
    {
        MSG_DEBUG(DEBUG_FFMPEG, "[ERROR] : [sendH264Frame] Could not allocate packet\n");
        return HI_FAILURE;
    }
    pkt->flags |= isKey? AV_PKT_FLAG_KEY : 0;
    pkt->size = size;
    pkt->data = data;
    pkt->stream_index = video_stream->index;

    /* input us time_stamp, trans to video_stream->timebase */
    if (h264_sps_pps_header.start_pts == 0 || h264_sps_pps_header.start_pts > pts_)
    {
        h264_sps_pps_header.start_pts = pts_;
    }
    
    pkt->pts = av_rescale_q(pts_ - h264_sps_pps_header.start_pts, us_time_base, video_stream->time_base);
    // pkt->pts = av_rescale_q(pts_idx, video_stream->time_base, video_stream->time_base);
    pkt->dts = pkt->pts;
	pkt->pos = -1;

    // pkt->pts = av_rescale_q(pts_, s_tTime_base, video_stream->time_base);
    // pkt->pts = av_rescale_q(h264_sps_pps_header.pts_idx, , video_stream->time_base);
	// pkt->dts = av_rescale_q_rnd(pkt->dts, s_tTime_base, video_stream->time_base, (AV_ROUND_NEAR_INF | AV_ROUND_PASS_MINMAX));
	// pkt->duration = av_rescale_q(pkt->duration, s_tTime_base, video_stream->time_base);
    // h264_sps_pps_header.pts_idx++;

    // printf("frame pkt %d >>>>>>>>>>\n", pkt->size);
    // for (int i = 0; i < 10; i++)
    // {
    //     printf("%0X ", pkt->data[i]);
    // }
    // printf("\n");

    pthread_mutex_lock(&(write_frame_mutex));
    s32Ret = av_interleaved_write_frame(output_format_context, pkt);
    if (s32Ret < 0)
    {
        MSG_DEBUG(DEBUG_FFMPEG, "[ERROR] : [writeH264PktForSend] send h264 packet fail.\n");
    }

    pthread_mutex_unlock(&(write_frame_mutex));
    av_packet_unref(pkt);
    av_packet_free(&pkt);
    return s32Ret;
}

HI_S32 hisi_tx_rtmp_video(const uint8_t *frame, int size, uint64_t ts)
{
    H264E_NALU_TYPE_E frame_type = 9;
    HI_BOOL isKeyFrame = HI_FALSE;
    if (*(frame) == 0x00 && *(frame + 1) == 0x00 && *(frame + 2) == 0x00 && *(frame + 3) == 0x01)
    {
        frame_type = *(frame + 4) & 0x1F;
        // if (frame_type == H264E_NALU_SEI) return HI_FAILURE;

        if (!h264_sps_pps_header.has_sps_pps)
        {
            if (frame_type == H264E_NALU_SPS)
            {
                memcpy(h264_sps_pps_header.sps_pps_buf, frame, size);
                h264_sps_pps_header.sps_pps_size = size;
            }
            if (frame_type == H264E_NALU_PPS && h264_sps_pps_header.sps_pps_size > 0)
            {
                memcpy(h264_sps_pps_header.sps_pps_buf + h264_sps_pps_header.sps_pps_size, frame, size);
                h264_sps_pps_header.sps_pps_size += size;
                h264_sps_pps_header.has_sps_pps = HI_TRUE;
            }
        }

        if (h264_sps_pps_header.has_sps_pps && !h264_sps_pps_header.isStreamInit)
        {
            video_stream->codecpar->extradata_size = h264_sps_pps_header.sps_pps_size;
            video_stream->codecpar->extradata = (uint8_t*)av_malloc(h264_sps_pps_header.sps_pps_size + AV_INPUT_BUFFER_PADDING_SIZE);
            if (video_stream->codecpar->extradata == NULL)
            {
                MSG_DEBUG(DEBUG_FFMPEG, "[ERROR] : [writeH264PktForSend] Could not av_malloc the video params extradata!\n");
                return HI_FAILURE;
            }
            memcpy(video_stream->codecpar->extradata, h264_sps_pps_header.sps_pps_buf, h264_sps_pps_header.sps_pps_size);

            /* step 4: open rtmp server connect and send header context */
            if (avio_open(&output_format_context->pb, rtmp_url, AVIO_FLAG_WRITE) < 0)
            {
                MSG_DEBUG(DEBUG_FFMPEG, "[ERROR] : [initFFmpegNetwork] open rtmp url failed.\n");
                return HI_FAILURE;
            }

            if (avformat_write_header(output_format_context, NULL) < 0)
            {
                MSG_DEBUG(DEBUG_FFMPEG, "[ERROR] : [initFFmpegNetwork] avformat_write_header failed.\n");
                return HI_FAILURE;
            }

            av_dump_format(output_format_context, 0, rtmp_url, 1);
            h264_sps_pps_header.isStreamInit = HI_TRUE;
            return HI_FAILURE;
        }

        if (!h264_sps_pps_header.isStreamInit || !h264_sps_pps_header.has_sps_pps) return HI_FAILURE; 

        /* step : send frame */
        if (frame_type == H264E_NALU_IDRSLICE) isKeyFrame = HI_TRUE;
        sendH264Frame(frame, size, isKeyFrame, ts);
    }
    return HI_SUCCESS;
}

static HI_S32 sendAacFrame(uint8_t* data, HI_U32 size, uint64_t pts_)
{
    HI_S32 s32Ret = HI_SUCCESS;

    /* step 1: init aac packet */
    AVPacket *pkt = av_packet_alloc();
    if (!pkt)
    {
        MSG_DEBUG(DEBUG_FFMPEG, "[ERROR] : [sendAacFrame] Could not allocate packet\n");
        return HI_FAILURE;
    }

    pkt->size = size;
    pkt->data = data;
    pkt->stream_index = audio_stream->index;

    /* input us time_stamp, trans to video_stream->timebase */
    if (h264_sps_pps_header.start_pts == 0 || h264_sps_pps_header.start_pts > pts_)
    {
        h264_sps_pps_header.start_pts = pts_;
    }
    
    pkt->pts = av_rescale_q(pts_ - h264_sps_pps_header.start_pts, us_time_base, audio_stream->time_base);
    pkt->dts = pkt->pts;
	pkt->pos = -1;

    pthread_mutex_lock(&(write_frame_mutex));
    s32Ret = av_interleaved_write_frame(output_format_context, pkt);
    if (s32Ret < 0)
    {
        MSG_DEBUG(DEBUG_FFMPEG, "[ERROR] : [sendAacFrame] send aac packet fail.\n");
    }

    pthread_mutex_unlock(&(write_frame_mutex));
    av_packet_unref(pkt);
    av_packet_free(&pkt);
    return s32Ret;
}

HI_S32 hisi_tx_rtmp_audio(const uint8_t *frame, int size, uint64_t ts)
{
    if (h264_sps_pps_header.isStreamInit != HI_TRUE) return HI_FAILURE;
    
    sendAacFrame(frame, size, ts);

    return HI_SUCCESS;
}